"""
电磁场数据加载和处理模块
"""

import torch
import numpy as np
import h5py
import os
import glob
from torch.utils.data import Dataset, DataLoader


class EMFieldDataset(Dataset):
    """电磁场数据集类"""
    
    def __init__(
        self,
        data_dir,
        field_names=['velocity', 'magnetic_field'],
        output_names='output-????',
        sub_x=1,
        sub_y=1,
        sub_z=1,
        sub_t=1,
        ind_x=None,
        ind_y=None,
        ind_z=None,
        ind_t=None,
        normalize=True,
        train=True
    ):
        """初始化电磁场数据集
        
        参数:
            data_dir: 数据目录
            field_names: 场名称列表
            output_names: 输出文件名模式
            sub_x, sub_y, sub_z: 空间采样间隔
            sub_t: 时间采样间隔
            ind_x, ind_y, ind_z: 空间索引上限
            ind_t: 时间索引上限
            normalize: 是否归一化数据
            train: 是否为训练集
        """
        self.data_dir = data_dir
        self.field_names = field_names
        self.output_names = output_names
        self.sub_x = sub_x
        self.sub_y = sub_y
        self.sub_z = sub_z
        self.sub_t = sub_t
        self.ind_x = ind_x
        self.ind_y = ind_y
        self.ind_z = ind_z
        self.ind_t = ind_t
        self.normalize = normalize
        self.train = train
        
        # 获取数据文件列表
        pattern = os.path.join(data_dir, output_names)
        self.file_list = sorted(glob.glob(pattern))
        
        # 读取第一个文件以获取元数据
        if len(self.file_list) > 0:
            self._load_metadata()
    
    def _load_metadata(self):
        """加载元数据"""
        with h5py.File(self.file_list[0], 'r') as f:
            # 获取坐标
            self.t = f['t'][:self.ind_t:self.sub_t]
            self.x = f['x'][:self.ind_x:self.sub_x]
            self.y = f['y'][:self.ind_y:self.sub_y]
            self.z = f['z'][:self.ind_z:self.sub_z]
            
            # 获取维度
            self.nt = len(self.t)
            self.nx = len(self.x)
            self.ny = len(self.y)
            self.nz = len(self.z)
            
            # 获取场形状
            field_shape = f[self.field_names[0]].shape
            self.n_fields = len(self.field_names)
            
            # 计算归一化参数
            if self.normalize:
                self._compute_normalization_params()
    
    def _compute_normalization_params(self):
        """计算归一化参数"""
        # 这里可以添加计算归一化参数的代码
        # 例如，计算所有数据文件的均值和标准差
        self.mean = np.zeros(4)  # u, v, Bx, By
        self.std = np.ones(4)    # u, v, Bx, By
    
    def __len__(self):
        """返回数据集大小"""
        return len(self.file_list)
    
    def __getitem__(self, idx):
        """获取单个数据样本
        
        参数:
            idx: 样本索引
            
        返回:
            inputs: 输入数据，形状为 [nt, nx, ny, nz, in_dim]
            outputs: 输出数据，形状为 [nt, nx, ny, nz, out_dim]
        """
        with h5py.File(self.file_list[idx], 'r') as f:
            # 读取场数据
            fields = {}
            for field_name in self.field_names:
                fields[field_name] = f[field_name][:]
            
            # 提取速度和磁场分量
            if 'velocity' in fields:
                velocity = fields['velocity']
                u = velocity[
                    :self.ind_t:self.sub_t,
                    :self.ind_x:self.sub_x,
                    :self.ind_y:self.sub_y,
                    :self.ind_z:self.sub_z,
                    0
                ]
                v = velocity[
                    :self.ind_t:self.sub_t,
                    :self.ind_x:self.sub_x,
                    :self.ind_y:self.sub_y,
                    :self.ind_z:self.sub_z,
                    1
                ]
                w = velocity[
                    :self.ind_t:self.sub_t,
                    :self.ind_x:self.sub_x,
                    :self.ind_y:self.sub_y,
                    :self.ind_z:self.sub_z,
                    2
                ]
            else:
                # 如果没有速度场，创建零场
                u = np.zeros((
                    self.nt, self.nx, self.ny, self.nz
                ))
                v = np.zeros((
                    self.nt, self.nx, self.ny, self.nz
                ))
                w = np.zeros((
                    self.nt, self.nx, self.ny, self.nz
                ))
            
            if 'magnetic_field' in fields:
                magnetic_field = fields['magnetic_field']
                Bx = magnetic_field[
                    :self.ind_t:self.sub_t,
                    :self.ind_x:self.sub_x,
                    :self.ind_y:self.sub_y,
                    :self.ind_z:self.sub_z,
                    0
                ]
                By = magnetic_field[
                    :self.ind_t:self.sub_t,
                    :self.ind_x:self.sub_x,
                    :self.ind_y:self.sub_y,
                    :self.ind_z:self.sub_z,
                    1
                ]
                Bz = magnetic_field[
                    :self.ind_t:self.sub_t,
                    :self.ind_x:self.sub_x,
                    :self.ind_y:self.sub_y,
                    :self.ind_z:self.sub_z,
                    2
                ]
            else:
                # 如果没有磁场，创建零场
                Bx = np.zeros((
                    self.nt, self.nx, self.ny, self.nz
                ))
                By = np.zeros((
                    self.nt, self.nx, self.ny, self.nz
                ))
                Bz = np.zeros((
                    self.nt, self.nx, self.ny, self.nz
                ))
            
            # 创建网格
            t_grid, x_grid, y_grid, z_grid = np.meshgrid(
                self.t, self.x, self.y, self.z, indexing='ij'
            )
            
            # 构建输入和输出
            # 输入: t, x, y, z, u0, v0, w0 (初始时刻的速度)
            # 输出: u, v, Bx, By (所有时刻的速度和磁场)
            data0 = np.stack([u[0], v[0], w[0]], axis=-1)  # 初始时刻的速度
            
            inputs = np.stack([
                t_grid, x_grid, y_grid, z_grid,
                np.repeat(data0[np.newaxis, ...], self.nt, axis=0)[..., 0],
                np.repeat(data0[np.newaxis, ...], self.nt, axis=0)[..., 1],
                np.repeat(data0[np.newaxis, ...], self.nt, axis=0)[..., 2]
            ], axis=-1)
            
            outputs = np.stack([u, v, Bx, By], axis=-1)
            
            # 归一化
            if self.normalize:
                inputs = (inputs - self.mean) / self.std
                outputs = (outputs - self.mean[:4]) / self.std[:4]
            
            # 转换为PyTorch张量
            inputs = torch.from_numpy(inputs).float()
            outputs = torch.from_numpy(outputs).float()
            
            return inputs, outputs


class EMFieldDataLoader:
    """电磁场数据加载器"""
    
    def __init__(
        self,
        dataset,
        batch_size=1,
        shuffle=True,
        num_workers=0,
        pin_memory=False,
        distributed=False
    ):
        """初始化数据加载器
        
        参数:
            dataset: 数据集
            batch_size: 批次大小
            shuffle: 是否打乱数据
            num_workers: 工作进程数
            pin_memory: 是否固定内存
            distributed: 是否分布式训练
        """
        self.dataset = dataset
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.num_workers = num_workers
        self.pin_memory = pin_memory
        self.distributed = distributed
        
        # 创建采样器
        if distributed:
            self.sampler = torch.utils.data.distributed.DistributedSampler(dataset)
        else:
            self.sampler = None
        
        # 创建数据加载器
        self.dataloader = DataLoader(
            dataset,
            batch_size=batch_size,
            shuffle=(shuffle and self.sampler is None),
            sampler=self.sampler,
            num_workers=num_workers,
            pin_memory=pin_memory,
            drop_last=True
        )
    
    def get_dataloader(self):
        """获取数据加载器"""
        return self.dataloader, self.sampler


class SyntheticEMDataGenerator:
    """合成电磁场数据生成器"""
    
    def __init__(
        self,
        nx=64,
        ny=64,
        nz=64,
        nt=50,
        Lx=1.0,
        Ly=1.0,
        Lz=1.0,
        T=1.0,
        device='cpu'
    ):
        """初始化合成数据生成器
        
        参数:
            nx, ny, nz: 空间网格点数
            nt: 时间步数
            Lx, Ly, Lz: 空间域大小
            T: 时间域大小
            device: 计算设备
        """
        self.nx = nx
        self.ny = ny
        self.nz = nz
        self.nt = nt
        self.Lx = Lx
        self.Ly = Ly
        self.Lz = Lz
        self.T = T
        self.device = device
        
        # 创建网格
        self.x = torch.linspace(0, Lx, nx, device=device)
        self.y = torch.linspace(0, Ly, ny, device=device)
        self.z = torch.linspace(0, Lz, nz, device=device)
        self.t = torch.linspace(0, T, nt, device=device)
        
        self.X, self.Y, self.Z, self.T = torch.meshgrid(
            self.x, self.y, self.z, self.t, indexing='ij'
        )
    
    def generate_mhd_data(self, num_samples=10):
        """生成MHD数据
        
        参数:
            num_samples: 样本数量
            
        返回:
            inputs: 输入数据，形状为 [num_samples, nt, nx, ny, nz, 7]
            outputs: 输出数据，形状为 [num_samples, nt, nx, ny, nz, 4]
        """
        inputs_list = []
        outputs_list = []
        
        for _ in range(num_samples):
            # 生成随机参数
            u0 = 0.1 * torch.randn(1, device=self.device)
            v0 = 0.1 * torch.randn(1, device=self.device)
            w0 = 0.1 * torch.randn(1, device=self.device)
            Bx0 = 0.1 * torch.randn(1, device=self.device)
            By0 = 0.1 * torch.randn(1, device=self.device)
            Bz0 = 0.1 * torch.randn(1, device=self.device)
            
            # 生成速度场
            u = u0 * torch.sin(2 * np.pi * self.X / self.Lx) * torch.cos(2 * np.pi * self.Y / self.Ly) * torch.cos(2 * np.pi * self.Z / self.Lz) * torch.cos(np.pi * self.T / self.T)
            v = v0 * torch.cos(2 * np.pi * self.X / self.Lx) * torch.sin(2 * np.pi * self.Y / self.Ly) * torch.cos(2 * np.pi * self.Z / self.Lz) * torch.cos(np.pi * self.T / self.T)
            w = w0 * torch.cos(2 * np.pi * self.X / self.Lx) * torch.cos(2 * np.pi * self.Y / self.Ly) * torch.sin(2 * np.pi * self.Z / self.Lz) * torch.cos(np.pi * self.T / self.T)
            
            # 生成磁场
            Bx = Bx0 * torch.cos(2 * np.pi * self.X / self.Lx) * torch.sin(2 * np.pi * self.Y / self.Ly) * torch.sin(2 * np.pi * self.Z / self.Lz) * torch.cos(np.pi * self.T / self.T)
            By = By0 * torch.sin(2 * np.pi * self.X / self.Lx) * torch.cos(2 * np.pi * self.Y / self.Ly) * torch.sin(2 * np.pi * self.Z / self.Lz) * torch.cos(np.pi * self.T / self.T)
            Bz = Bz0 * torch.sin(2 * np.pi * self.X / self.Lx) * torch.sin(2 * np.pi * self.Y / self.Ly) * torch.cos(2 * np.pi * self.Z / self.Lz) * torch.cos(np.pi * self.T / self.T)
            
            # 构建输入和输出
            # 输入: t, x, y, z, u0, v0, w0 (初始时刻的速度)
            # 输出: u, v, Bx, By (所有时刻的速度和磁场)
            data0 = torch.stack([u[0], v[0], w[0]], dim=-1)  # 初始时刻的速度
            
            inputs = torch.stack([
                self.T, self.X, self.Y, self.Z,
                torch.repeat_interleave(data0[..., 0:1, np.newaxis, np.newaxis, np.newaxis], self.nt, dim=0).squeeze(-1),
                torch.repeat_interleave(data0[..., 1:2, np.newaxis, np.newaxis, np.newaxis], self.nt, dim=0).squeeze(-1),
                torch.repeat_interleave(data0[..., 2:3, np.newaxis, np.newaxis, np.newaxis], self.nt, dim=0).squeeze(-1)
            ], dim=-1)
            
            outputs = torch.stack([u, v, Bx, By], dim=-1)
            
            inputs_list.append(inputs)
            outputs_list.append(outputs)
        
        # 堆叠所有样本
        inputs = torch.stack(inputs_list, dim=0)
        outputs = torch.stack(outputs_list, dim=0)
        
        return inputs, outputs
    
    def generate_maxwell_data(self, num_samples=10):
        """生成麦克斯韦方程数据
        
        参数:
            num_samples: 样本数量
            
        返回:
            inputs: 输入数据，形状为 [num_samples, nt, nx, ny, nz, 7]
            outputs: 输出数据，形状为 [num_samples, nt, nx, ny, nz, 6]
        """
        inputs_list = []
        outputs_list = []
        
        for _ in range(num_samples):
            # 生成随机参数
            E0 = 0.1 * torch.randn(3, device=self.device)
            B0 = 0.1 * torch.randn(3, device=self.device)
            
            # 生成电场
            Ex = E0[0] * torch.sin(2 * np.pi * self.X / self.Lx) * torch.cos(2 * np.pi * self.Y / self.Ly) * torch.cos(2 * np.pi * self.Z / self.Lz) * torch.cos(np.pi * self.T / self.T)
            Ey = E0[1] * torch.cos(2 * np.pi * self.X / self.Lx) * torch.sin(2 * np.pi * self.Y / self.Ly) * torch.cos(2 * np.pi * self.Z / self.Lz) * torch.cos(np.pi * self.T / self.T)
            Ez = E0[2] * torch.cos(2 * np.pi * self.X / self.Lx) * torch.cos(2 * np.pi * self.Y / self.Ly) * torch.sin(2 * np.pi * self.Z / self.Lz) * torch.cos(np.pi * self.T / self.T)
            
            # 生成磁场
            Bx = B0[0] * torch.cos(2 * np.pi * self.X / self.Lx) * torch.sin(2 * np.pi * self.Y / self.Ly) * torch.sin(2 * np.pi * self.Z / self.Lz) * torch.cos(np.pi * self.T / self.T)
            By = B0[1] * torch.sin(2 * np.pi * self.X / self.Lx) * torch.cos(2 * np.pi * self.Y / self.Ly) * torch.sin(2 * np.pi * self.Z / self.Lz) * torch.cos(np.pi * self.T / self.T)
            Bz = B0[2] * torch.sin(2 * np.pi * self.X / self.Lx) * torch.sin(2 * np.pi * self.Y / self.Ly) * torch.cos(2 * np.pi * self.Z / self.Lz) * torch.cos(np.pi * self.T / self.T)
            
            # 构建输入和输出
            # 输入: t, x, y, z, Ex0, Ey0, Ez0 (初始时刻的电场)
            # 输出: Ex, Ey, Ez, Bx, By, Bz (所有时刻的电场和磁场)
            data0 = torch.stack([Ex[0], Ey[0], Ez[0]], dim=-1)  # 初始时刻的电场
            
            inputs = torch.stack([
                self.T, self.X, self.Y, self.Z,
                torch.repeat_interleave(data0[..., 0:1, np.newaxis, np.newaxis, np.newaxis], self.nt, dim=0).squeeze(-1),
                torch.repeat_interleave(data0[..., 1:2, np.newaxis, np.newaxis, np.newaxis], self.nt, dim=0).squeeze(-1),
                torch.repeat_interleave(data0[..., 2:3, np.newaxis, np.newaxis, np.newaxis], self.nt, dim=0).squeeze(-1)
            ], dim=-1)
            
            outputs = torch.stack([Ex, Ey, Ez, Bx, By, Bz], dim=-1)
            
            inputs_list.append(inputs)
            outputs_list.append(outputs)
        
        # 堆叠所有样本
        inputs = torch.stack(inputs_list, dim=0)
        outputs = torch.stack(outputs_list, dim=0)
        
        return inputs, outputs
    
    def save_data(self, inputs, outputs, save_dir):
        """保存数据到文件
        
        参数:
            inputs: 输入数据
            outputs: 输出数据
            save_dir: 保存目录
        """
        os.makedirs(save_dir, exist_ok=True)
        
        for i in range(inputs.shape[0]):
            with h5py.File(os.path.join(save_dir, f'output-{i:04d}.h5'), 'w') as f:
                # 保存坐标
                f.create_dataset('t', data=self.t.cpu().numpy())
                f.create_dataset('x', data=self.x.cpu().numpy())
                f.create_dataset('y', data=self.y.cpu().numpy())
                f.create_dataset('z', data=self.z.cpu().numpy())
                
                # 保存场数据
                if outputs.shape[-1] == 4:  # MHD数据
                    velocity = np.zeros((self.nt, self.nx, self.ny, self.nz, 3))
                    velocity[..., 0] = inputs[i, ..., 4].cpu().numpy()  # u
                    velocity[..., 1] = inputs[i, ..., 5].cpu().numpy()  # v
                    velocity[..., 2] = np.zeros((self.nt, self.nx, self.ny, self.nz))  # w
                    
                    magnetic_field = np.zeros((self.nt, self.nx, self.ny, self.nz, 3))
                    magnetic_field[..., 0] = outputs[i, ..., 2].cpu().numpy()  # Bx
                    magnetic_field[..., 1] = outputs[i, ..., 3].cpu().numpy()  # By
                    magnetic_field[..., 2] = np.zeros((self.nt, self.nx, self.ny, self.nz))  # Bz
                    
                    f.create_dataset('velocity', data=velocity)
                    f.create_dataset('magnetic_field', data=magnetic_field)
                
                elif outputs.shape[-1] == 6:  # 麦克斯韦数据
                    electric_field = np.zeros((self.nt, self.nx, self.ny, self.nz, 3))
                    electric_field[..., 0] = outputs[i, ..., 0].cpu().numpy()  # Ex
                    electric_field[..., 1] = outputs[i, ..., 1].cpu().numpy()  # Ey
                    electric_field[..., 2] = outputs[i, ..., 2].cpu().numpy()  # Ez
                    
                    magnetic_field = np.zeros((self.nt, self.nx, self.ny, self.nz, 3))
                    magnetic_field[..., 0] = outputs[i, ..., 3].cpu().numpy()  # Bx
                    magnetic_field[..., 1] = outputs[i, ..., 4].cpu().numpy()  # By
                    magnetic_field[..., 2] = outputs[i, ..., 5].cpu().numpy()  # Bz
                    
                    f.create_dataset('electric_field', data=electric_field)
                    f.create_dataset('magnetic_field', data=magnetic_field)